import torchvision.models as models
from torchvision import transforms
from torch.autograd import Variable
import os, torch
import torch.nn as nn

from src.Fer2013DataSet import Fer2013DataSet


class Res18Feature(nn.Module):
    def __init__(self, pretrained, num_classes=10):
        super(Res18Feature, self).__init__()
        resnet = models.resnet18(pretrained)
        # self.feature = nn.Sequential(*list(resnet.children())[:-2]) # before avgpool
        self.features = nn.Sequential(*list(resnet.children())[:-1])  # after avgpool 512x1

        fc_in_dim = list(resnet.children())[-1].in_features  # original fc layer's in dimention 512

        self.fc = nn.Linear(fc_in_dim, num_classes)  # new fc layer 512x7
        self.alpha = nn.Sequential(nn.Linear(fc_in_dim, 1), nn.Sigmoid())

    def forward(self, x):
        x = self.features(x)

        x = x.view(x.size(0), -1)

        attention_weights = self.alpha(x)
        out = attention_weights * self.fc(x)
        return attention_weights, out


# 模型存储路径
model_save_path = "src/models/new_epoch35_acc0.8221.pth"  # 修改为你自己保存下来的模型文件
img_path = "test_3066_aligned.jpg"  # 待测试照片位置

# ------------------------ 加载数据 --------------------------- #

res18 = Res18Feature(pretrained=False)
checkpoint = torch.load(model_save_path)
res18.load_state_dict(checkpoint['model_state_dict'])
res18.cuda(0)
res18.eval()
data_transforms_val = transforms.Compose([
    transforms.ToPILImage(),
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0.225])])
val_dataset = Fer2013DataSet("/home/zbzbzzz/datasets/fer2013/img", phase='valid', transform=data_transforms_val)
correct_sum = 0
all_sum = 0
for i in range(7000):
    all_sum += 1
    img, label, idx = val_dataset.__getitem__(i)
    tensor = Variable(torch.unsqueeze(img, dim=0).float(), requires_grad=False)
    tensor = tensor.cuda(0)
    _, outputs = res18(tensor)
    _, predicts = torch.max(outputs, 1)
    print("======================")
    if predicts == label:
        print("正确")
        correct_sum += 1
    else:
        print("错误")
    print(predicts, label)
acc = correct_sum / all_sum * 100
print("正确率：", acc)
# 正确率： 17.014613778705638
# model_save_path = "src/models/epoch50_acc0.827.pth"
